#include "rclcpp/rclcpp.hpp"
#include "geometry_msgs/msg/twist.hpp"
#include "turtlesim/msg/pose.hpp"
#include "rclcpp/utilities.hpp"
#include <cmath>

using namespace std::chrono_literals;

class SquareMoveNode : public rclcpp::Node
{
public:
    SquareMoveNode() : Node("square_move_node")
    {
        this->declare_parameter<double>("side_length", 2.0);
        side_length_ = this->get_parameter("side_length").as_double();

        publisher_ = this->create_publisher<geometry_msgs::msg::Twist>("/turtle1/cmd_vel", 10);
        
        pose_subscriber_ = this->create_subscription<turtlesim::msg::Pose>(
            "/turtle1/pose", 10, std::bind(&SquareMoveNode::pose_callback, this, std::placeholders::_1));

        timer_ = this->create_wall_timer(20ms, std::bind(&SquareMoveNode::timer_callback, this)); // 提高到20ms

        state_ = MOVE_FORWARD;
        move_start_x_ = 0.0;
        move_start_y_ = 0.0;
        linear_speed_ = 1.0;
        
        // PID控制参数
        kp_ = 6.0;  // 比例系数
        ki_ = 0.5;  // 积分系数  
        kd_ = 0.1;  // 微分系数
        
        // 初始化
        first_pose_received_ = false;
        turn_count_ = 0;
        integral_error_ = 0.0;
        prev_error_ = 0.0;

        RCLCPP_INFO(this->get_logger(), "高精度正方形节点已启动，边长: %.2f", side_length_);
    }

private:
    enum State{
        MOVE_FORWARD,
        TURN
    };

    rclcpp::Publisher<geometry_msgs::msg::Twist>::SharedPtr publisher_;
    rclcpp::Subscription<turtlesim::msg::Pose>::SharedPtr pose_subscriber_;
    rclcpp::TimerBase::SharedPtr timer_;
    State state_;
    double side_length_;
    double linear_speed_;
    
    // PID控制参数
    double kp_, ki_, kd_;
    double integral_error_;
    double prev_error_;
    
    // 位置相关变量
    double current_x_;
    double current_y_;
    double current_angle_;
    double move_start_x_;
    double move_start_y_;
    double turn_start_angle_;
    double target_angle_;
    bool first_pose_received_;
    int turn_count_;

    void pose_callback(const turtlesim::msg::Pose::SharedPtr msg)
    {
        current_x_ = msg->x;
        current_y_ = msg->y;
        current_angle_ = msg->theta;
        
        if (!first_pose_received_) {
            first_pose_received_ = true;
            move_start_x_ = current_x_;
            move_start_y_ = current_y_;
        }
    }

    // 规范化角度到 [-π, π] 范围
    double normalize_angle(double angle)
    {
        angle = std::fmod(angle, 2 * M_PI);
        if (angle > M_PI) angle -= 2 * M_PI;
        if (angle < -M_PI) angle += 2 * M_PI;
        return angle;
    }

    // PID控制器
    double pid_control(double error, double dt)
    {
        // 积分项（抗饱和）
        integral_error_ += error * dt;
        // 限制积分项防止windup
        if (integral_error_ > 1.0) integral_error_ = 1.0;
        if (integral_error_ < -1.0) integral_error_ = -1.0;
        
        // 微分项
        double derivative = (error - prev_error_) / dt;
        prev_error_ = error;
        
        // PID输出
        double output = kp_ * error + ki_ * integral_error_ + kd_ * derivative;
        
        // 输出限幅
        if (output > 3.0) output = 3.0;
        if (output < -3.0) output = -3.0;
        
        return output;
    }

    void timer_callback()
    {
        if (!first_pose_received_) {
            return;
        }

        auto message = geometry_msgs::msg::Twist();
        static auto last_time = this->now();
        auto current_time = this->now();
        double dt = (current_time - last_time).seconds();
        last_time = current_time;

        switch(state_){
            case MOVE_FORWARD:
                {
                    // 重置PID控制器
                    integral_error_ = 0.0;
                    prev_error_ = 0.0;
                    
                    message.linear.x = linear_speed_;
                    message.angular.z = 0.0;
                    
                    // 计算从起点移动的实际距离
                    double dx = current_x_ - move_start_x_;
                    double dy = current_y_ - move_start_y_;
                    double actual_distance = std::sqrt(dx * dx + dy * dy);
                    
                    // 在最后10cm减速
                    double distance_to_go = side_length_ - actual_distance;
                    if (distance_to_go < 0.1) {
                        message.linear.x = linear_speed_ * (distance_to_go / 0.1);
                        if (message.linear.x < 0.1) message.linear.x = 0.1;
                    }
                    
                    // 检查是否到达目标距离
                    if(actual_distance >= side_length_ - 0.01){ // 留1cm余量
                        state_ = TURN;
                        turn_start_angle_ = current_angle_;
                        target_angle_ = normalize_angle(turn_start_angle_ + M_PI / 2);
                        
                        turn_count_++;
                        RCLCPP_INFO(this->get_logger(), 
                                   "第%d次转弯: 当前角度 %.2f°, 目标角度 %.2f°", 
                                   turn_count_,
                                   turn_start_angle_ * 180 / M_PI, 
                                   target_angle_ * 180 / M_PI);
                    }
                }
                break;

            case TURN:
                {
                    message.linear.x = 0.0;
                    
                    // 计算角度差
                    double angle_diff = normalize_angle(target_angle_ - current_angle_);
                    
                    // 使用PID控制
                    double angular_control = pid_control(angle_diff, dt);
                    message.angular.z = angular_control;
                    
                    // 检查是否达到目标角度（非常严格的误差）
                    if (std::abs(angle_diff) < 0.005) { // 约0.3度的误差
                        state_ = MOVE_FORWARD;
                        move_start_x_ = current_x_;
                        move_start_y_ = current_y_;
                        
                        RCLCPP_INFO(this->get_logger(), 
                                   "第%d次转弯完成! 最终角度 %.2f°, 目标角度 %.2f°, 误差 %.2f°", 
                                   turn_count_,
                                   current_angle_ * 180 / M_PI,
                                   target_angle_ * 180 / M_PI,
                                   std::abs(angle_diff) * 180 / M_PI);
                        
                        if (turn_count_ >= 4) {
                            RCLCPP_INFO(this->get_logger(), "完美正方形完成!");
                        }
                    } else if (std::abs(angle_diff) < 0.05) { // 约3度时开始详细记录
                        RCLCPP_DEBUG(this->get_logger(), 
                                    "精调中: 当前 %.3f°, 目标 %.3f°, 误差 %.3f°", 
                                    current_angle_ * 180 / M_PI,
                                    target_angle_ * 180 / M_PI,
                                    angle_diff * 180 / M_PI);
                    }
                }
                break;
        }

        publisher_->publish(message);
    }
};

int main(int argc, char * argv[])
{
    rclcpp::init(argc, argv);
    auto node = std::make_shared<SquareMoveNode>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}